# -*- coding: utf-8 -*-
import random


class Corpus:
    """语料数据类，封装了根据索引获取问题和回答、语料采样方法
    
    Attributes
    ----------
    questions : 数组， 各种问题
    answers : 数组， 与questions一一对应
    """

    def __init__(self, questions, answers):
        self._questions = questions
        self._answers = answers

    def get_question(self, index):
        """根据索引获取问题
        Args
        ----
        index : 整数或者整数列表， 当索引为整数列表时，返回的questions也是列表

        Returns
        -------
        ques : str or list of str
        """
        if isinstance(index, int):
            return self._questions[index]
        elif isinstance(index, list):
            return [self.get_question(ind) for ind in index]
        else:
            raise ValueError("index must be an int number or a list of int")

    def get_answer(self, index):
        """根据索引获取对应的答案
        Args
        ----
        index : 整数或者整数列表， 当索引为整数列表时，返回的回答也是列表

        Returns
        -------
        ques : str or list of str
        """
        if isinstance(index, int):
            return self._answers[index]
        elif isinstance(index, list):
            return [self.get_answer(ind) for ind in index]
        else:
            raise ValueError("index must be an int number or a list of int")

    def take_samples(self, N, question=None):
        """抽取N个样本， 返回这些样本的索引列表
        Args
        ----
        ask : 字符串， 问题相关的ask (该方法的一个可能实现逻辑是从语料中得到N个与ask相近的样本)
        N ： 样本的数量

        Returns
        -------
        sub_indexes :数组，样本的索引
        """
        questions_len = len(self._questions)
        questions_index = range(questions_len)
        samples_index = random.sample(questions_index, min(questions_len, N))
        return samples_index

